通过 --use_distributed_optimizer 来开启基于分布式优化器的 ZeRO-1 优化。Stage 1 的优化对 SGD 没有效果，但是依然在 GPT-2 训练脚本中加了相关 option，为的是之后 stage 2/3 的扩展。

实现细节：

1. 添加 ParamAndGradBucket/Group/Buffer 的基建，采用了与 Megatron-LM 类似的实现思路：所有 param/grad 都连续放在一个一维 buffer 中，并分 bucket 按组进行通信；
2. 修改 DistributedDataParallel 的逻辑，需要在构造函数中创建上述 buffer、划分 bucket group，并完成 hook 注册；同时保留了原 reducer 逻辑的实现分支，不开启 --use_distributed_optimizer 时则走原先的逻辑；
3. 添加 DistributedOptimizer 类，继承自 Optimizer；同时为每种优化器包了一个仅接受 models->Parameters() 的构造 function 方法，DistributedOptimizer 在构造时会接受这个 function，并在其构造函数内部创建 base_optimizer；
4. 修改了 PP 的构造函数参数，把 optimizer 的传入延后到了 Module::TrainStep，这样 PP 对象可以完全不持有 optimizer 对象；
5. 在训练循环添加了本轮峰值占用/预留显存信息的输出。

详细实现流程可以参考叙述：https://gxtctab8no8.feishu.cn/wiki/XQbGwXSsZi3MutkZuhXcKWsinnY#share-QU0fdM1cYoT06vxmwMbcnLJUn1b

两个比较重要的细节：

1. DDP 对象在构造时，会一同构造模型对应的 buffer/bucket_group 等；同时 DistOpt 在构造时会接受 DDP 对象的 buffer/group；
2. 由于现在 PP + DDP 的实现，会单独将每个 chunk 构造为 DDP 对象，所以传给 DistOpt 的 buffer/bucket_group 等，需要额外在 chunk 之间做一道汇总，得到一个总 list 再传进去。